from cgitb import text
from genericpath import isfile
import json
import os
from random import Random
import numpy
from typing import List, Tuple


path = './data/matadata/'

def remove_substring(sentence:str, istart:int, iend:int, labels:List[str]) -> Tuple[str, List[str]]:
    iend = min(iend, len(sentence))
    if istart == iend:
        return sentence
    before = sentence[:istart]
    after = sentence[iend:]
    before_labels = labels[:istart]
    after_labels = labels[iend:]
    return before + after, before_labels + after_labels

def remove_substring_till_next_section(sentence:str, istart:int, labels:List[str]) -> Tuple[str, List[str]]:
    for j in range(istart + 1, len(sentence)):
        if sentence[j] == '【':
            return remove_substring(sentence, istart, j, labels)
    return remove_substring(sentence, istart, 99999, labels)

unused_section_titles = []
for line in open('./data/matadata/todelete.txt', mode='r', encoding='utf-8-sig'):
    line = line.strip()
    unused_section_titles.append(line)


sentences = []
for line in open('./data/matadata/text.txt', mode='r', encoding='utf-8-sig'):
    line = line.strip()
    sentences.append(line)
infos = []
for line in open('./data/matadata/labels.txt', mode='r', encoding='utf-8-sig'):
    line = line.strip()
    infos.append(line)
assert(len(sentences) == len(infos))
sentences_with_labels:List[Tuple[str, List[str]]] = []
for i, (sentence, info) in enumerate(zip(sentences, infos)):

    if '发现躯干前倾、侧凸进行性加重伴腰痛' in sentence:
        kkk = 0

    sentence_parallel_labels = ['O'] * len(sentence)
    if info != '':
        for label_info in json.loads(info):
            istart = label_info['start']
            iend = label_info['end']
            label_text = label_info['text']
            label_type = label_info['labels'][0]
            #substr = sentence[istart:iend]
            #assert( substr == label_text)
            sentence_parallel_labels[istart] = 'B-' + label_type
            sentence_parallel_labels[istart + 1 : iend] = ['I-' + label_type] * (iend - istart - 1)
    sentences_with_labels.append((sentence, sentence_parallel_labels))
new_sentence_with_labels:List[Tuple[str, List[str]]] = []
for sentence, labels in sentences_with_labels:
    for title in unused_section_titles:
        i = sentence.find(title)
        if i != -1:
            sentence, labels = remove_substring_till_next_section(sentence, i, labels)
        else:
            continue
    new_sentence_with_labels.append((sentence, labels))

r = Random()
r.shuffle(new_sentence_with_labels)
split_pos = round(len(new_sentence_with_labels) * 0.7)
train_set = new_sentence_with_labels[:split_pos]
test_set = new_sentence_with_labels[split_pos:]
with open('./data/matadata/train_plain.txt', mode='w', encoding='utf8') as f:
    for sentence, labels in train_set:
        if sentence != '':
            f.write(f'{sentence}\n{" ".join(labels)}\n')
with open('./data/matadata/test_plain.txt', mode='w', encoding='utf8') as f:
    for sentence, labels in test_set:
        if sentence != '':
            f.write(f'{sentence}\n{" ".join(labels)}\n')